import random

from util.common import *

if __name__ == '__main__':

    list = []

    df = read_excel("source/tf_sa_zx_vuln_info.xlsx")
    filtered_df = df[df['VULN_ID'].notna()]

    df2 = read_excel("source/tf_smp_web_asset2.xlsx")
    filtered_df2 = df2[df2['INTRANET_IP'].notna()]

    for createTime in generate_workdays('2024-09-01', '2024-11-28', 5):

        timeSpan = random.randrange(2, 20)
        start_date, end_date = get_random_date(createTime, timeSpan)

        num_samples = random.randint(1, 3)
        random_vuln = filtered_df.sample(n=num_samples)
        random_assert = filtered_df2.sample(n=1)

        for row in random_vuln.itertuples():

            assetId = random_assert['ID'].item()
            iIP = random_assert["INTRANET_IP"].item()
            iPort = random_assert["INTRANET_PORT"].item()
            ip6 = random_assert["INTRANET_IP_VAL"].item()
            url = random_assert["INTRANET_URL"].item()
            assetName = random_assert["ASSET_NAME"].item()

            id = "VULNZX" + random_string(14)
            vid = row.VULN_ID

            seep = {
                "id": id,
                "vid": vid,
                "assetId": assetId,
                "ip": iIP,
                "ip6": ip6,
                "port": iPort,
                "url": url,
                "assetName": assetName,
                "startTime": start_date,
                "endTime": end_date,
                "creator": 'PORT6452766823027939',
            }
            message1 = render_template("""
               INSERT INTO tf_sa_zx_vuln_asset(REL_ID, VULN_ID, TARGET_TYPE, ASSET_ID, IP, IP_VAL, PORT, URL, ASSET_NAME,
                                ASSET_TYPE_ID, ASSET_OS_ID, ASSET_GROUP_ID, IS_REPAIR, REPAIR_TIME, REPAIR_MEASURES,
                                CHECK_RESULT, CREATE_TIME, UPDATE_TIME)
                VALUES ($id, $vid, '1', $assetId, $ip,
                        $ip6, $port, $url, $assetName, '1', NULL, '0',
                        NULL, NULL, NULL, NULL, $startTime,
                        $endTime);
            """, seep)

            list.append(message1)

            report = {
                "id": id,
                "vid": vid,
                "taskId": id,
                "assetId": assetId,
                "ip": iIP,
                "ip6": ip6,
                "port": iPort,
                "url": url,
                "assetName": assetName,
                "startTime": start_date,
                "endTime": end_date,
                "creator": 'PORT6452766823027939',
            }

            message2 = render_template(
                """
                       INSERT INTO tf_sa_zx_vuln_task(TASK_ID, TASK_FLAG, TASK_STATUS, TASK_MAIN_ID, WO_ID, VULN_ID, ACTUAL_START_TIME,
                               ACTUAL_END_TIME, EXE_TYPE, EXE_PARAM, CREATOR, CREATE_TIME, UPDATOR, UPDATE_TIME)
                        VALUES ($id, '2', '5', NULL, NULL, $vid,
                                $startTime,
                                $endTime, '1',
                                '{"exeStartTime":"","exeEndTime":"","exeType":"1","executeCron":""}', $creator,
                                $startTime, $creator,
                                $endTime);
                        
                        INSERT INTO tf_sa_zx_vuln_task_rel(REL_ID, TASK_ID, TARGET_TYPE, ASSET_ID, IP, IP_VAL, PORT, URL, ASSET_NAME,
                                                           ASSET_TYPE_ID, ASSET_OS_ID, ASSET_GROUP_ID, CHECK_STATUS, FAIL_REASON)
                        VALUES ($id, $taskId, '1', $assetId, $ip,
                                $ip6, $port, $url, $assetName, '1',
                                NULL, '0', '0', NULL);
                        
                        INSERT INTO tf_sa_zx_vuln_task_result(RESULT_ID, TASK_ID, ASSET_ID, VULN_ID, CHECK_RESULT, CREATE_TIME, UPDATE_TIME)
                        VALUES ($id, $taskId, $assetId, $vid, '0',
                                $startTime,
                                $endTime);
                    """,
                report)
            list.append(message2)

    write_file('target/SATP_专项漏洞管理造数.sql', list)
       
